# pytorch_linear_regression_example.py
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt

# ========== 1. 生成示例数据 ==========
# y = 2x + 1 + 一点噪声
torch.manual_seed(0)
X = torch.unsqueeze(torch.linspace(-5, 5, 100), dim=1)  # shape: (100, 1)
y = 2 * X + 1 + 0.5 * torch.randn(X.size())  # 添加随机噪声

# ========== 2. 定义模型 ==========
class LinearModel(nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        self.linear = nn.Linear(1, 1)  # 输入1维，输出1维

    def forward(self, x):
        return self.linear(x)

model = LinearModel()

# ========== 3. 定义损失函数与优化器 ==========
criterion = nn.MSELoss()           # 均方误差
optimizer = optim.SGD(model.parameters(), lr=0.01)  # 随机梯度下降

# ========== 4. 训练模型 ==========
epochs = 200
for epoch in range(epochs):
    # 前向传播
    outputs = model(X)
    loss = criterion(outputs, y)
    
    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if (epoch+1) % 20 == 0:
        print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}")

# ========== 5. 可视化训练结果 ==========
predicted = model(X).detach()
plt.scatter(X.numpy(), y.numpy(), label='True Data')
plt.plot(X.numpy(), predicted.numpy(), color='r', label='Fitted Line')
plt.legend()
plt.show()

# ========== 6. 保存模型 ==========
torch.save(model.state_dict(), "linear_model.pt")
print("模型已保存为 linear_model.pt")

# ========== 7. 加载模型并预测新数据 ==========
loaded_model = LinearModel()
loaded_model.load_state_dict(torch.load("linear_model.pt"))
loaded_model.eval()  # 切换到评估模式

# 预测新的数据点
new_X = torch.tensor([[4.0], [10.0], [-3.5]])  # 示例：3个新数据点
preds = loaded_model(new_X)
print("\n新的数据点:")
print(new_X)
print("预测结果:")
print(preds)
